package degrade

import (
	"errors"
	"sync"
	"time"
)

/**
令牌桶算法


算法思想
想象有一个木桶，以固定的速度往木桶里加入令牌，木桶满了则不再加入令牌。服务收到请求时尝试从木桶中取出一个令牌，如果能够得到令牌则继续执行后续的业务逻辑；
如果没有得到令牌，直接返回反问频率超限的错误码或页面等，不继续执行后续的业务逻辑

特点：由于木桶内只要有令牌，请求就可以被处理，所以令牌桶算法可以支持突发流量。同时由于往木桶添加令牌的速度是固定的，
且木桶的容量有上限，所以单位时间内处理的请求书也能够得到控制，起到限流的目的。假设加入令牌的速度为 1token/10ms，桶的容量为500，
在请求比较的少的时候（小于每10毫秒1个请求）时，木桶可以先"攒"一些令牌（最多500个）。当有突发流量时，一下把木桶内的令牌取空，
也就是有500个在并发执行的业务逻辑，之后要等每10ms补充一个新的令牌才能接收一个新的请求。

参数设置：木桶的容量 - 考虑业务逻辑的资源消耗和机器能承载并发处理多少业务逻辑。生成令牌的速度 - 太慢的话起不到“攒”令牌应对突发流量的效果。

适用场景：
适合电商抢购或者微博出现热点事件这种场景，因为在限流的同时可以应对一定的突发流量。如果采用均匀速度处理请求的算法，在发生热点时间的时候，
会造成大量的用户无法访问，对用户体验的损害比较大。

go语言实现：
假设每100ms生产一个令牌，按user_id/IP记录访问最近一次访问的时间戳 t_last 和令牌数，每次请求时如果 now - last > 100ms, 增加 (now - last) / 100ms个令牌。
然后，如果令牌数 > 0，令牌数 -1 继续执行后续的业务逻辑，否则返回请求频率超限的错误码或页面。
*/

const (
	METHOD                = "METHOD"
	_METHOD_DEFAULT       = "DEFAULT"
	_THOUSAND             = 1000
	_MILLION        int64 = 1e6
)

/*
 *valid 本令牌桶对象是否有效
 *rate 生产令牌的速率，每rate微秒生产一个令牌
 *capacity 令牌桶容量
 *stored 令牌桶存量
 *nextProducePoint 下一次生产令牌的时间点
 */
type tokenBucket struct {
	valid            bool
	rate             float64 //token  生产速率
	capacity         float64 //桶容量
	stored           float64
	nextProducePoint int64
	gLock            *sync.RWMutex
}

func NewTokenBucket(cap float64, rate float64) *tokenBucket {
	return &tokenBucket{
		valid:            true,
		rate:             rate,
		capacity:         cap,
		stored:           0.0,
		nextProducePoint: 0,
		gLock:            &sync.RWMutex{},
	}
}

func (p *tokenBucket) clear() {
	p.valid = false
	p.rate = 0.0
	p.capacity = 0.0
	p.stored = 0.0
	p.nextProducePoint = 0
}

/*
 * 设置限流rate(qps)
 *
 * @return
 */
func (p *tokenBucket) setRate(rate float64) {
	if rate <= 0.0 {
		p.valid = false
		return
	}
	p.gLock.Lock()
	p.rate = float64(_MILLION) / rate
	p.gLock.Unlock()
}

/*
 * 获取限流rate(qps)
 *
 * @return float64  rate(qps)
 */
func (p *tokenBucket) getRate() float64 {
	p.gLock.RLock()
	defer p.gLock.RUnlock()
	return float64(_MILLION) / p.rate
}

/*
 * 获取桶大小
 *
 * @return float64  capacity
 */
func (p *tokenBucket) getCapacity() float64 {
	return p.capacity
}

/*
 * 获取令牌
 * @param[in] float64 : 希望获取的令牌数
 * @param[in] int64 : 可以等待的时间，单位ms
 *
 * @return nil:获取成功; error:获取失败
 */
func (p *tokenBucket) IsLimited() bool {
	if err := p.tryAcquire(1, 0); err != nil {
		return false
	}
	return true
}

func (p *tokenBucket) tryAcquire(permits float64, timeout int64) error {
	if p.valid == false {
		return errors.New("TokenBucket invalid")
	}
	p.gLock.Lock()
	//时间单位统一为微秒
	now := time.Now().UnixNano() / _THOUSAND
	if p.stored < permits && p.nextProducePoint > now+timeout*_THOUSAND {
		p.gLock.Unlock()
		return errors.New("Acquire failed")
	}
	wait := p.acquire(permits)
	p.gLock.Unlock()
	if wait > 0 && timeout > 0 {
		time.Sleep(time.Duration(wait) * time.Microsecond)
	}
	return nil
}

/*
 * 同步到now时间点的状态
 */
func (p *tokenBucket) sync(now int64) {
	if now > p.nextProducePoint {
		nextTimeStored := p.stored + float64(now-p.nextProducePoint)/p.rate
		if p.capacity <= nextTimeStored {
			p.stored = p.capacity
		} else {
			p.stored = nextTimeStored
		}
		p.nextProducePoint = now
	}
}

/*
 * 获取令牌
 */
func (p *tokenBucket) acquire(permits float64) int64 {
	now := time.Now().UnixNano() / _THOUSAND
	p.sync(now)
	//可以拿的量
	var getPermits float64
	if permits <= p.stored {
		getPermits = permits
	} else {
		getPermits = p.stored
	}
	//除了库存还需要拿的量
	stillNeed := permits - getPermits
	//如果拿了超出库存，那么需要更新下一生产点
	p.nextProducePoint += int64(stillNeed * p.rate)

	p.stored -= getPermits
	return p.nextProducePoint - now
}

/*
 *name 本节点名称
 *keys 可以命中本节点的搜索key
 *quotaProportion 本节点配额(比例)
 *quota 本节点总配额
 *burst 本节点burst（其实叶子节点的burst才有效）
 *level 本节点层级
 *config 下层map
 */
type quotaConfig struct {
	name            string
	keys            map[string]int
	quotaProportion float64
	quota           float64
	burst           float64
	config          map[string]*quotaConfig
}
